'''
https://leetcode.cn/problems/01-matrix/description/?envType=study-plan-v2&envId=graph-theory
'''
from collections import deque
from typing import List

class Solution:
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        m, n = len(mat), len(mat[0])
        direction = [(1, 0), (-1, 0), (0, 1), (0, -1)]
        q = deque()
        visited = [[False] * n for _ in range(m)]
        dist = [[0] * n for _ in range(m)]
        for i in range(m):
            for j in range(n):
                if mat[i][j] == 0:
                    visited[i][j] = True
                    q.append((i, j))
        step = 0
        while q:
            for _ in range(len(q)):
                i, j = q.popleft()
                for a,b in direction:
                    tx, ty = i + a, j + b
                    if 0 <= tx < m and 0 <= ty < n and not visited[tx][ty] and mat[tx][ty] == 1:
                        dist[tx][ty] = step + 1
                        visited[tx][ty] = True
                        q.append((tx, ty))
            step += 1
        return dist

mat = [[0, 0, 0], [0, 1, 0], [1, 1, 1]]
print(Solution().updateMatrix(mat))